package primary.code08;

import lombok.extern.slf4j.Slf4j;

import java.util.*;

/**
 * 并查集数据结构
 * 快速地实现了两种操作：
 * 1. 将两个集合合并成一个集合
 * 2. 判断两个点是否处于一个集合中
 * <p>
 * 实现：
 * 将用户数据封装成“上指图的节点”，根节点是所在集合的代表。
 * 操作1：将小集合的代表节点指向大集合代表节点；
 * 操作2：向上查询代表节点，代表节点相同则表示处于同一集合，否则反之；
 */
@Slf4j
public class Code03_UnionFindSet<T> {

    static class Element<T> {
        public T val;

        public Element(T val) {
            this.val = val;
        }
    }

    private final Map<T, Element<T>> elementMap = new HashMap<>();

    private final Map<Element, Element> fatherMap = new HashMap<>();

    private final Map<Element, Integer> sizeMap = new HashMap<>();

    public Code03_UnionFindSet(List<T> vals) {
        for (T val : vals) {
            Element<T> e = new Element<>(val);
            elementMap.put(val, e);
            fatherMap.put(e, e);
            sizeMap.put(e, 1);
        }
    }

    public void union(T v1, T v2) {
        Element<T> n1 = elementMap.get(v1);
        Element<T> n2 = elementMap.get(v2);
        if (n1 == null || n2 == null) {
            log.warn("this val is not in set, cannot union");
            return;
        }
        Element<T> represent1 = findRepresent(n1);
        Element<T> represent2 = findRepresent(n2);
        boolean isSameSet = represent1 == represent2;
        if (isSameSet) {
            log.info("already in same set");
            return;
        }
        Element<T> small = sizeMap.get(represent1) < sizeMap.get(represent2) ? represent1 : represent2;
        Element<T> big = small == represent1 ? represent2 : represent1;
        fatherMap.put(small, big);
        sizeMap.put(big, sizeMap.get(big) + sizeMap.get(small));
        sizeMap.remove(small);
    }

    public boolean isSameSet(T v1, T v2) {
        Element<T> n1 = elementMap.get(v1);
        Element<T> n2 = elementMap.get(v2);
        if (n1 == null || n2 == null) {
            return false;
        }
        Element<T> represent1 = findRepresent(n1);
        Element<T> represent2 = findRepresent(n2);
        return represent1 == represent2;
    }

    private Element<T> findRepresent(Element<T> node) {
        Stack<Element<T>> path = new Stack<>(); // 记录路径
        Element<T> cur = node;
        Element<T> father = fatherMap.get(cur);
        path.push(cur);
        if (father != node) {
            cur = father;
            path.push(cur);
            father = fatherMap.get(cur);
        }
        while (!path.isEmpty()) { // 将路径上的节点直接指向代表节点
            Element<T> next = path.pop();
            fatherMap.put(next, father);
        }
        return father;
    }


    public static void main(String[] args) {
        Code03_UnionFindSet unionFindSet = new Code03_UnionFindSet(Arrays.asList(1, 2, 3, 4, 5, 6, 7, 8, 9));
        unionFindSet.union(0, 2);
        unionFindSet.union(1, 2);
        unionFindSet.union(2, 3);

        unionFindSet.union(8, 9);
        unionFindSet.union(6, 8);

        unionFindSet.union(3, 9);

        log.info("{} and {} in same set? {}", 1, 3, unionFindSet.isSameSet(0, 3)); //false
        log.info("{} and {} in same set? {}", 1, 3, unionFindSet.isSameSet(1, 3)); //true
        log.info("{} and {} in same set? {}", 8, 9, unionFindSet.isSameSet(8, 9)); //true
        log.info("{} and {} in same set? {}", 8, 9, unionFindSet.isSameSet(2, 6)); //true
        log.info("{} and {} in same set? {}", 1, 6, unionFindSet.isSameSet(1, 6)); //true
        log.info("{} and {} in same set? {}", 4, 6, unionFindSet.isSameSet(4, 6)); //false
    }
}
